{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fda134cd-247d-40e4-b37f-0ab0a2d57943",
   "metadata": {},
   "source": [
    "# LangGraph with AgentCore Memory Tool (Short term memory)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98765432-abcd-1234-efgh-ijklmnopqrst",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "This notebook demonstrates how to integrate Amazon Bedrock AgentCore Memory capabilities with a conversational AI agent using LangGraph framework. We'll focus on **short-term memory** retention within a single conversation session - allowing an agent to recall information from earlier in the conversation without explicit context management.\n",
    "\n",
    "\n",
    "## Tutorial Details\n",
    "\n",
    "| Information         | Details                                                                          |\n",
    "|:--------------------|:---------------------------------------------------------------------------------|\n",
    "| Tutorial type       | Short Term Conversational                                                        |\n",
    "| Agent usecase       | Personal Fitness                                                                 |\n",
    "| Agentic Framework   | Langgraph                                                                        |\n",
    "| LLM model           | Anthropic Claude Sonnet 3                                                        |\n",
    "| Tutorial components | AgentCore Short-term Memory, Langgraph, Memory retrieval via Tool                |\n",
    "| Example complexity  | Beginner                                                                         |\n",
    "\n",
    "You'll learn to:\n",
    "- Create a memory store with AgentCore Memory for short-term memory\n",
    "- Use LangGraph to create an agent with structured memory workflows\n",
    "- Implement memory tools for conversation history retrieval\n",
    "- Access and utilize contextual information within a single session\n",
    "- Enhance conversational experiences through effective memory recall\n",
    "\n",
    "\n",
    "### Scenario Context\n",
    "\n",
    "In this example, we'll create a \"**Personal Fitness Coach**\" that can remember workout details, fitness goals, physical limitations, and exercise preferences as they are mentioned throughout the conversation. This assistant will demonstrate how effective short-term memory management enables a more natural and personalized fitness coaching experience without requiring users to repeatedly state their information.\n",
    "\n",
    "\n",
    "## Architecture\n",
    "<div style=\"text-align:left\">\n",
    "    <img src=\"architecture.png\" width=\"65%\" />\n",
    "</div>\n",
    "\n",
    "## Prerequisites\n",
    "\n",
    "- Python 3.10+\n",
    "- AWS account with appropriate permissions\n",
    "- AWS IAM role with appropriate permissions for AgentCore Memory\n",
    "- Access to Amazon Bedrock models\n",
    "\n",
    "Let's get started by setting up our environment!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d6fc488-0f70-416b-a0e9-57671e0bfc7d",
   "metadata": {},
   "source": [
    "## Step 1: Environment Setup\n",
    "Let's begin importing all the necessary libraries and defining the clients to make this notebook work."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36e53844-5916-4470-a656-d59c64af4e84",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -qr requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa41d38e-859c-4634-a5ef-624cbe6e63ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "from datetime import datetime"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "849b4372-ac00-49a8-bcdf-916ac895b10a",
   "metadata": {},
   "source": [
    "Define the region and the role with the appropiate permissions for Amazon Bedrock models and AgentCore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f366c0b-b410-4fe5-8935-c330d5ce04d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "region = os.getenv('AWS_REGION', 'us-west-2')\n",
    "\n",
    "logging.basicConfig(level=logging.INFO, format=\"%(asctime)s - %(levelname)s - %(message)s\", datefmt=\"%Y-%m-%d %H:%M:%S\")\n",
    "logger = logging.getLogger(\"agentcore-memory\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1b2c3d4-e5f6-a7b8-c9d0-e1f2a3b4c5d6",
   "metadata": {},
   "source": [
    "### How the Integration Works\n",
    "\n",
    "The integration between LangGraph and AgentCore Memory involves:\n",
    "\n",
    "1. Using AgentCore Memory to store conversations in the short term memory\n",
    "2. Structured workflows in LangGraph to manage memory operations\n",
    "\n",
    "This approach separates memory management from reasoning, creating a cleaner and more maintainable agent architecture."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f68f7df8-2904-43a7-8dfb-fdaa72e2f88e",
   "metadata": {},
   "source": [
    "## Step 2: Memory Creation\n",
    "In this section, we'll create a memory store using the AgentCore Memory SDK. This memory store will allow our agent to retain information from the conversation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04d31ed2-78db-42c1-9392-8815229a4a27",
   "metadata": {},
   "outputs": [],
   "source": [
    "from bedrock_agentcore.memory import MemoryClient\n",
    "from botocore.exceptions import ClientError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff6ea0be-5073-4990-a2b4-3c44510da3bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "client = MemoryClient(region_name=region)\n",
    "memory_name = \"FitnessCoach\"\n",
    "memory_id = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18b2c138-b452-4b2a-9730-6e41c824ee94",
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    print(\"Creating Memory...\")\n",
    "    # Create the memory resource\n",
    "    memory = client.create_memory_and_wait(\n",
    "        name=memory_name,                       # This name is unique across all memories in this account\n",
    "        description=\"Fitness Coach Agent\",      # Human-readable description\n",
    "        strategies=[],                          # No memory strategies for short-term memory\n",
    "        event_expiry_days=7,                    # Memories expire after 7 days\n",
    "        max_wait=300,                           # Maximum time to wait for memory creation (5 minutes)\n",
    "        poll_interval=10                        # Check status every 10 seconds\n",
    "    )\n",
    "\n",
    "    # Extract and print the memory ID\n",
    "    memory_id = memory['id']\n",
    "    logger.info(f\"Memory created successfully with ID: {memory_id}\")\n",
    "except ClientError as e:\n",
    "    if e.response['Error']['Code'] == 'ValidationException' and \"already exists\" in str(e):\n",
    "        # If memory already exists, retrieve its ID\n",
    "        memories = client.list_memories()\n",
    "        memory_id = next((m['id'] for m in memories if m['id'].startswith(memory_name)), None)\n",
    "        logger.info(f\"Memory already exists. Using existing memory ID: {memory_id}\")\n",
    "except Exception as e:\n",
    "    # Handle any errors during memory creation\n",
    "    logger.info(f\"❌ ERROR: {e}\")\n",
    "    import traceback\n",
    "    traceback.print_exc()\n",
    "    # Cleanup on error - delete the memory if it was partially created\n",
    "    if memory_id:\n",
    "        try:\n",
    "            client.delete_memory_and_wait(memory_id=memory_id)\n",
    "            logger.info(f\"Cleaned up memory: {memory_id}\")\n",
    "        except Exception as cleanup_error:\n",
    "            logger.info(f\"Failed to clean up memory: {cleanup_error}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7fac6fd5-6228-4a20-88c0-ae30da48cf3f",
   "metadata": {},
   "source": [
    "## Step 3: LangGraph Agent Creation\n",
    "Let's import all the libraries we need to create the agent with LangGraph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86d2e879-0697-4157-9877-3253b135716a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.graph import StateGraph, MessagesState\n",
    "from langgraph.prebuilt import ToolNode, tools_condition\n",
    "from langchain_core.tools import tool\n",
    "from langchain_core.messages import HumanMessage, SystemMessage\n",
    "from langchain_aws import ChatBedrock"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10c0914c-6af5-4242-9dec-b3f1a05a6432",
   "metadata": {},
   "source": [
    "### LangGraph Agent Implementation\n",
    "\n",
    "Now let's create the agent with LangGraph, incorporating our memory tools:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0fa8672a-f5b6-43b6-b2fe-5fedc080a090",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_agent(client, memory_id, actor_id, session_id):\n",
    "    \"\"\"Create and configure the LangGraph agent\"\"\"\n",
    "    \n",
    "    # Initialize your LLM (adjust model and parameters as needed)\n",
    "    llm = ChatBedrock(\n",
    "        model_id=\"anthropic.claude-3-sonnet-20240229-v1:0\",  # or your preferred model\n",
    "        model_kwargs={\"temperature\": 0.1}\n",
    "    )\n",
    "    \n",
    "    @tool\n",
    "    def list_events():\n",
    "        \"\"\"Tool used when needed to retrieve recent information\"\"\" \n",
    "        events = client.list_events(\n",
    "                memory_id=memory_id,\n",
    "                actor_id=actor_id,\n",
    "                session_id=session_id,\n",
    "                max_results=10\n",
    "            )\n",
    "        return events\n",
    "        \n",
    "    \n",
    "    # Bind tools to the LLM\n",
    "    tools = [list_events]\n",
    "    llm_with_tools = llm.bind_tools(tools)\n",
    "    \n",
    "    # System message\n",
    "    system_message = \"\"\"You are the Personal Fitness Coach, a sophisticated fitness guidance assistant.\n",
    "                        PURPOSE:\n",
    "                        - Help users develop workout routines based on their fitness goals\n",
    "                        - Remember user's exercise preferences, limitations, and progress\n",
    "                        - Provide personalized fitness recommendations and training plans\n",
    "                        MEMORY CAPABILITIES:\n",
    "                        - You have access to recent events with the list_events tool\n",
    "                        \"\"\"\n",
    "    \n",
    "    # Define the chatbot node\n",
    "    def chatbot(state: MessagesState):\n",
    "        raw_messages = state[\"messages\"]\n",
    "    \n",
    "        # Remove any existing system messages to avoid duplicates or misplacement\n",
    "        non_system_messages = [msg for msg in raw_messages if not isinstance(msg, SystemMessage)]\n",
    "    \n",
    "        # Always ensure SystemMessage is first\n",
    "        messages = [SystemMessage(content=system_message)] + non_system_messages\n",
    "    \n",
    "        latest_user_message = next((msg.content for msg in reversed(messages) if isinstance(msg, HumanMessage)), None)\n",
    "    \n",
    "        # Get response from model with tools bound\n",
    "        response = llm_with_tools.invoke(messages)\n",
    "    \n",
    "        # Save conversation if applicable\n",
    "        if latest_user_message and response.content.strip():  # Check that response has content\n",
    "            conversation = [\n",
    "                (latest_user_message, \"USER\"),\n",
    "                (response.content, \"ASSISTANT\")\n",
    "            ]\n",
    "            \n",
    "            # Validate that all message texts are non-empty\n",
    "            if all(msg[0].strip() for msg in conversation):  # Ensure no empty messages\n",
    "                try:\n",
    "                    client.create_event(\n",
    "                        memory_id=memory_id,\n",
    "                        actor_id=actor_id,\n",
    "                        session_id=session_id,\n",
    "                        messages=conversation\n",
    "                    )\n",
    "                except Exception as e:\n",
    "                    print(f\"Error saving conversation: {str(e)}\")\n",
    "        \n",
    "        # Append response to full message history\n",
    "        return {\"messages\": raw_messages + [response]}\n",
    "    \n",
    "    # Create the graph\n",
    "    graph_builder = StateGraph(MessagesState)\n",
    "    \n",
    "    # Add nodes\n",
    "    graph_builder.add_node(\"chatbot\", chatbot)\n",
    "    graph_builder.add_node(\"tools\", ToolNode(tools))\n",
    "    \n",
    "    # Add edges\n",
    "    graph_builder.add_conditional_edges(\n",
    "        \"chatbot\",\n",
    "        tools_condition,\n",
    "    )\n",
    "    graph_builder.add_edge(\"tools\", \"chatbot\")\n",
    "    \n",
    "    # Set entry point\n",
    "    graph_builder.set_entry_point(\"chatbot\")\n",
    "    \n",
    "    # Compile the graph\n",
    "    return graph_builder.compile()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "929386fc-542f-49c9-9541-daeac8f32348",
   "metadata": {},
   "source": [
    "### Creating a Wrapper for Agent Invocation\n",
    "\n",
    "Let's create a simple wrapper to invoke our agent:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "979f87e0-a547-448b-9ef9-b416351c9916",
   "metadata": {},
   "outputs": [],
   "source": [
    "def langgraph_bedrock(payload, agent):\n",
    "    \"\"\"\n",
    "    Invoke the agent with a payload\n",
    "    \"\"\"\n",
    "    user_input = payload.get(\"prompt\")\n",
    "    \n",
    "    # Create the input in the format expected by LangGraph\n",
    "    response = agent.invoke({\"messages\": [HumanMessage(content=user_input)]})\n",
    "    \n",
    "    # Extract the final message content\n",
    "    return response[\"messages\"][-1].content"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43528c42-384a-4822-970b-929875d50757",
   "metadata": {},
   "source": [
    "## Step 4: Run the LangGraph Agent\n",
    "We can now run the agent with our AgentCore Memory integration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2beeaaea-5e84-483a-aea5-a31e7e2b6769",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create unique actor and session IDs for this conversation\n",
    "actor_id = f\"user-{datetime.now().strftime('%Y%m%d%H%M%S')}\"\n",
    "session_id = f\"workout-{datetime.now().strftime('%Y%m%d%H%M%S')}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09d2576d-88da-4b7c-bbe4-34f2bf97db7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the agent with AgentCore Memory integration\n",
    "agent = create_agent(client, memory_id, actor_id, session_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22334455-6677-8899-aabb-ccddeeff0011",
   "metadata": {},
   "source": [
    "#### Congratulations ! Your Agent is ready !!\n",
    "\n",
    "### Let's test the Agent\n",
    "\n",
    "Let's interact with our agent to test its memory capabilities:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db5530d8-a22b-4086-85ec-63435bed9219",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = langgraph_bedrock({\"prompt\": \"Hello! This is my first day, I need a workout routine.\"}, agent)\n",
    "print(f\"Agent: {response}\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aabbccdd-eeff-0011-2233-445566778899",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = langgraph_bedrock({\"prompt\": \"I want to build muscle, looking for a biceps routine. I have some lower back problems.\"}, agent)\n",
    "print(f\"Agent: {response}\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9988776655-4433-2211-0099-8877665544",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = langgraph_bedrock({\"prompt\": \"Can you give me three exercises with number of reps?\"}, agent)\n",
    "print(f\"Agent: {response}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a2b3c4d-5e6f-7g8h-9i0j-1k2l3m4n5o6p",
   "metadata": {},
   "source": [
    "### Testing Memory Persistence\n",
    "\n",
    "To truly demonstrate the power of the AgentCore Memory integration, let's create a new agent instance and see if it can recall our previous conversation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "q1w2e3r4t5y6-u7i8o9p0-a1s2d3f4-g5h6j7k8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a new agent instance (simulating a new session)\n",
    "new_agent = create_agent(client, memory_id, actor_id, session_id)\n",
    "\n",
    "# Test if the new agent remembers our preferences\n",
    "response = langgraph_bedrock({\n",
    "    \"prompt\": \"Hello again! Can you remind me about my last workout session?\"\n",
    "}, new_agent)\n",
    "\n",
    "print(\"New Agent Session:\\n\")\n",
    "print(f\"Agent: {response}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "z1x2c3v4b5n6m7-8k9j0h1g2-f3d4s5a6-p7o8i9u0",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "In this notebook, we've demonstrated:\n",
    "\n",
    "1. How to create a AgentCore Memory resource for an AI agent\n",
    "2. Building a LangGraph workflow with memory integration\n",
    "3. Implementing memory tools for conversation history retrieval\n",
    "4. Creating an agent that intelligently uses memory when needed\n",
    "5. Testing memory persistence across agent instances\n",
    "\n",
    "This integration showcases the power of combining structured workflows (LangGraph) with robust memory systems (AgentCore Memory) to create more intelligent and context-aware AI agents.\n",
    "\n",
    "The approach we've demonstrated can be extended to more complex use cases, including multi-agent systems, long-term memory with extraction strategies, and specialized memory retrieval based on conversation context."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "213d2a15-3af9-4a07-94b6-f7d39e0fa038",
   "metadata": {},
   "source": [
    "## Clean up\n",
    "Let's delete the memory to clean up the resources used in this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df3893ef-f6c3-4744-9860-ff15e014943c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#client.delete_memory_and_wait(memory_id = memory_id, max_wait = 300, poll_interval =10)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_pytorch_p310",
   "language": "python",
   "name": "conda_pytorch_p310"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
